import numpy as np
import pandas as pd
import json
import cv2


# 注意：请在images的同级目录下创建一个aug_small_images文件夹


#root_path = r'C:\Users\lbk-pc\Desktop\my_project\datas'
root_path = r'D:\datasets\chongqing'
train_path = root_path + '\\chongqing1_round1_train1_20191223'
train_img_path = train_path + '\\images\\'
train_json_file = train_path + '\\clean_annotations.json'

with open(train_json_file) as f:
    annot_json = json.load(f)
annot_json.keys()


# 存下类型字典备用
cat_names = {c['id']: c['name'] for c in annot_json['categories']}



def add_annot(train_path, images_list, annotations_list, categories_list):
    '''

    '''
    small_annotations['images'] = images_list
    small_annotations['annotations'] = annotations_list
    small_annotations['categories'] = categories_list

    small_cut_file = train_path + '\\small_annotations.json'
    with  open(small_cut_file, 'w')  as f:
        json.dump(small_annotations, f)


# 定义一些存annotations文件信息的变量
df_annotations = pd.DataFrame(annot_json['annotations'])
small_annotations = {}  #
images_list = []
annotations_list = []
new_name_number = 0

# n = 0  # 处理前2张图
# annot = annot_json['annotations'][64:]

print("start...")
for ann in annot_json['annotations']:
    ann_area = ann['area']
    if ann_area < 1024:
        # 计算截取框
        images_to_cut = {}  # 创建一个新图片的images
        annotations_to_cut = {}  # 创建一个新图片的annotations
        # print("start...")
        for ori_image in annot_json['images']:
            if ori_image['id'] == ann['image_id']:  # 根据image_id信息索引 h,w

                ori_image_name = ori_image['file_name']
                ori_image_id = ori_image['id']

                ## cut_xxx 在原图中要截取的区域
                cut_w = int(ori_image['width'] / 4)
                cut_h = int(ori_image['height'] / 4)
                cut_x = int(ann['bbox'][0] - (cut_w - ann['bbox'][2]) / 2)
                cut_y = int(ann['bbox'][1] - (cut_h - ann['bbox'][3]) / 2)

                if cut_x < 0:  # 处理左边x超出情况
                    cut_x = 0
                if cut_y < 0:  # 处理上边y超出情况
                    cut_y = 0
                if cut_x > (ori_image['width'] - cut_w):  # 处理右边x超出情况
                    cut_x = ori_image['width'] - cut_w
                if cut_y > (ori_image['height'] - cut_h):  # 处理下边y超出情况
                    cut_y = ori_image['height'] - cut_h

                # 获取新images信息
                new_img_name = "aug_img_" + str(new_name_number) + "_from_" + ori_image_name
                images_to_cut["file_name"] = new_img_name
                images_to_cut["height"] = cut_h
                images_to_cut["width"] = cut_w
                # images_to_cut["id"] = image_id
                images_to_cut["id"] = 10000 + new_name_number  # 新的id从10000开始
                images_list.append(images_to_cut)

                # 获取新categories信息
                categories_list = annot_json['categories']

                # 处理相同image_id包含多个瑕疵框信息
                number_bbox = 0
                df_annot = df_annotations[df_annotations['image_id'] == ori_image_id]
                for index, row in df_annot.iterrows():
                    if row['area'] < 1024:
                        # 左上角的点在截取框内 且 右下角的点在截取框内
                        if row['bbox'][0] > cut_x \
                                and row['bbox'][1] > cut_y  \
                                and  row['bbox'][0] + row['bbox'][2] < cut_x + cut_w \
                                and  row['bbox'][1] + row['bbox'][3] < cut_y + cut_h:

                            # 获取新annotations信息
                            annotations_to_cut = {}  # 创建一个新图片的annotations
                            annotations_to_cut = row.to_dict().copy()
                            new_xx = row['bbox'][0] - cut_x  # 同image_id图里的新坐标
                            new_yy = row['bbox'][1] - cut_y
                            new_ww = row['bbox'][2]
                            new_hh = row['bbox'][3]
                            annotations_to_cut['bbox'] = [new_xx, new_yy, new_ww, new_hh]
                            annotations_to_cut['image_id'] = 10000 + new_name_number  # 新的image_id从10000开始
                            annotations_list.append(annotations_to_cut)
                            number_bbox += 1


                # 剪裁图片
                img = cv2.imread(train_img_path + "\\" + ori_image_name)
                print('image_id:', ori_image_id)

                new_img = img[cut_y:cut_y + cut_h, cut_x:cut_x + cut_w]
                print('剪裁后图片维度',new_img.shape[0:2])

                print('剪裁后图片包含bbox数量', number_bbox)

                # n += 1
                new_name_number += 1
                print("已处理{}张小尺寸图片".format(new_name_number))
                print('-'*30)

                # 保存剪裁好的图片
                cv2.imwrite(train_path + "\\" + "aug_small_images" + "\\" + new_img_name, new_img)

    # if n == 2:  # 处理2张图片
    #     break

print("总共有{}个annotations信息".format(len(annotations_list)))
print("end...")

# 写入json文件
add_annot(train_path, images_list, annotations_list, categories_list)